from suffix_trees import STree
import time
from needle import Needle
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Align import MultipleSeqAlignment
from Bio.Seq import Seq
from Bio import AlignIO
import csv
import sys


class Solution:
    dataset = 'sars-cov2.fasta'
    genes = dict()  # 存放30个基因序列的字典
    _needle = Needle()
    alignments = []
    msgs = []

    def __init__(self):
        # self.s1_pos = []    # 不想要这么多列表
        # self.s2_pos = []
        # self.length = []
        self.pos1_pos2_len = []

    def get_genes(self):  # 很不错的解析方式
        for seq_record in SeqIO.parse(self.dataset, 'fasta'):
            self.genes[seq_record.id] = self.replace_special(str(seq_record.seq))

    @staticmethod
    def replace_special(line):
        return line.replace('R', 'AG').replace('Y', 'CT').replace('M', 'AC').replace('K', 'GT').replace('S',
                                                                                                        'CG').replace(
            'W', 'AT').replace('H', 'ACT').replace('B', 'CGT').replace('V', 'ACG').replace('D', 'AGT').replace('N',
                                                                                                               'ACGT')

    @staticmethod
    def read_file(file):  # 很菜的方式,已被丢弃
        with open(file, 'r', encoding='utf-8') as f:
            title = f.readline().strip().replace('>', '')
            string = ''
            for line in f:
                string += line.strip()
        return title, string

    @staticmethod
    def select_prefix(s2, index, tree_s1):
        length = 1
        _length = 1
        position = -1
        start = index  # 起始位置
        while index < len(s2) - 1:
            index += 1
            dest = tree_s1.find(s2[start:index + 1])

            if dest == -1:
                if position != -1:  # 修正非最远匹配
                    if _length != 1:
                        length = _length
                    return position, length
                else:
                    return -1, 0
            else:
                if abs(dest - start) < length + 1:  # 避免非最远匹配
                    position = dest
                    _length = length + 1
                length += 1  # useful
        return position, length

    # 查找公共子串
    def common_substring(self, s2, tree_s1):
        index = 0
        while index < len(s2) - 1:
            position, length = self.select_prefix(s2, index, tree_s1)
            if position != -1:
                if abs(position - index) < length:  # 要有交叠
                    # self.s1_pos.append(position)
                    # self.s2_pos.append(index)
                    # self.length.append(length)
                    self.pos1_pos2_len.append((position, index, length))
                    # print(index, position, length)
                    index += length
                else:
                    index += 1
            else:
                index += 1

    def greedy_common_list(self):  # 使用贪婪的方式,避免公共子串之间交叠
        sorted_lists = sorted(self.pos1_pos2_len, key=lambda l: l[0])
        valid_list = []
        cur_val = 0
        for x in sorted_lists:
            if x[0] >= cur_val:
                valid_list.append(x)
                cur_val = (x[0] + x[2])  # 出坑
        return valid_list

    def main(self):  # 主函数
        self.get_genes()
        index = 1
        genes_keys = self.genes.keys()
        for key1 in genes_keys:
            seq1 = self.genes[key1]
            tree = STree.STree(seq1)
            for key2 in genes_keys:
                if key1 == key2:
                    continue
                self.pos1_pos2_len.clear()  # 清空
                # self.s1_pos.clear()
                # self.s2_pos.clear()
                # self.length.clear()
                seq2 = self.genes[key2]
                start_time = time.time()
                self.common_substring(seq2, tree)
                valid_list = self.greedy_common_list()  # 最佳子串
                align1, align2, score = self.segment_align(valid_list, seq1, seq2)
                time_use = time.time() - start_time
                # 打印一些提示信息
                print('=' * 50)
                print(index)
                memory = self.get_size(tree)  # 查看树的内存
                print(key1, key2, sep='\t')
                print(len(seq1), len(seq2), sep='\t')
                self.get_summary(key1, align1, key2, align2, score, time_use, memory)
                if index % 100 == 0:
                    self.save_alignments()
                    self.alignments.clear()  # 避免内存占用
                index += 1
        del tree  # 避免占用
        self.save_summary()

    def get_size(self, obj, seen=None):
        """Recursively finds size of objects"""
        size = sys.getsizeof(obj)
        if seen is None:
            seen = set()
        obj_id = id(obj)
        if obj_id in seen:
            return 0
        # Important mark as seen *before* entering recursion to gracefully handle
        # self-referential objects
        seen.add(obj_id)
        if isinstance(obj, dict):
            size += sum([self.get_size(v, seen) for v in obj.values()])
            size += sum([self.get_size(k, seen) for k in obj.keys()])
        elif hasattr(obj, '__dict__'):
            size += self.get_size(obj.__dict__, seen)
        elif hasattr(obj, '__iter__') and not isinstance(obj, (str, bytes, bytearray)):
            size += sum([self.get_size(i, seen) for i in obj])
        return size

    def get_summary(self, key1, align1, key2, align2, score, _time, memory):  # 这真的是一件麻烦的事情
        msg = dict()
        msg['two genes'] = key1 + '_' + key2
        msg['align len'] = len(align1)
        msg['score'] = int(score)
        msg['time use'] = round(_time, 2)
        msg['suffix tree memory'] = memory
        self.msgs.append(msg)

        # 写入对齐文件
        _alignment = MultipleSeqAlignment([
            SeqRecord(Seq(align1), id=key1, description='|' + str(msg)),
            SeqRecord(Seq(align2), id=key2, description='|' + str(msg))
        ])
        self.alignments.append(_alignment)
        print(msg)

    def save_summary(self):
        with open('summary.csv', 'w', encoding='utf-8-sig', newline='') as f:
            fieldnames = self.msgs[0].keys()
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            for msg in self.msgs:
                writer.writerow(msg)

    def save_alignments(self):
        with open('alignments.fasta', 'a+', encoding='utf-8') as f:
            AlignIO.write(self.alignments, f, 'fasta')

    def segment_align(self, valid_list, seq1, seq2):
        align1 = ''
        align2 = ''  # 存放结果
        score = 0
        # 处理0序列
        _align1, _align2, _score = self._needle.needle(seq1[0:valid_list[0][0]], seq2[0:valid_list[0][1]])
        score += _score
        # score += self._needle.finalize(_align1, _align2)    # 精度更高,但我不想使用
        align1 += _align1
        align2 += _align2
        for i in range(len(valid_list) - 1):
            common_seq = seq1[valid_list[i][0]:valid_list[i][0] + valid_list[i][2]]
            align1 += common_seq
            align2 += common_seq
            score += len(common_seq) * 10  # 相同子串
            _align1, _align2, _score = self._needle.needle(
                seq1[valid_list[i][0] + valid_list[i][2]:valid_list[i + 1][0]],
                # 中间部分
                seq2[valid_list[i][1] + valid_list[i][2]:valid_list[i + 1][1]])
            score += _score
            align1 += _align1
            align2 += _align2
            # score += self._needle.finalize(_align1, _align2)
        # 处理最后的
        common_seq = seq1[valid_list[-1][0]:valid_list[-1][0] + valid_list[-1][2]]
        align1 += common_seq
        align2 += common_seq
        score += len(common_seq) * 10
        _align1, _align2, _score = self._needle.needle(seq1[valid_list[-1][0] + valid_list[-1][2]:len(seq1)],
                                                       seq2[valid_list[-1][1] + valid_list[-1][2]:len(seq2)])
        score += _score
        # score += self._needle.finalize(_align1, _align2)
        align1 += _align1
        align2 += _align2
        return align1, align2, score  # perfect done !


"""
if __name__ == '__main__':
    solution = Solution()
    key1, seq1 = solution.read_file('datasets/corynebacterium_diphtheriae_nctc.fasta')
    key2, seq2 = solution.read_file('datasets/corynebacterium_efficiens_ys-314.fasta')
    # seq1 = 'acgtacgatcgacgatacgtagctgatcaaaaa'
    # seq2 = 'gctagcatgatcgacgtcaaaaaaaaaaaa'
    # seq1 = 'gattcctaatcggcgacaaaagacatgcagttctccacagacaactacctggtgccgaag'
    # seq2 = 'ggagggtcgacgggatggcagagcgggtcgccggcgatccggagctcgccatcttatggt'
    print(len(seq1), len(seq2))
    tree = STree.STree(seq1)
    # print(tree.find('ggt'))
    # print(len(set1))
    # print(len(set2))
    # print(tree.find('ta'))
    # print(set2.find('ta'))
    solution.common_substring(seq2, tree)
    print(solution.s1_pos, solution.s2_pos, solution.length, sep='\n')
    print(solution.greedy_common_list())
"""

if __name__ == '__main__':
    solution = Solution()
    solution.main()

"""
if __name__ == '__main__':
    solution = Solution()
    solution.get_genes()
    print(solution.genes)
"""
